import os
import pandas as pd
import numpy as np
from fastparquet import ParquetFile

# Parquet File Path
PARQUET_PATH = os.path.join(os.path.dirname(__file__), "job/data", "trade_indicators.parquet")

def load_parquet():
    """Load and cache Parquet file"""
    print(f"Checking for parquet file at: {PARQUET_PATH}")
    if not os.path.exists(PARQUET_PATH):
        raise FileNotFoundError(f"股票数据文件不存在: {PARQUET_PATH}")

    # Fix _ARRAY_API error
    import fastparquet.schema
    if not hasattr(fastparquet.schema, '_ARRAY_API'):
        fastparquet.schema._ARRAY_API = np

    pf = ParquetFile(PARQUET_PATH)
    df = pf.to_pandas()
    df = df.reset_index()

    # Log available columns for debugging
    print(f"Available columns in Parquet file after reset_index: {list(df.columns)}")

    # Check for 'date' or 'datekey' column
    if 'date' in df.columns:
        # Ensure 'date' is in YYYY-MM-DD string format
        df['date'] = pd.to_datetime(df['date'], errors='coerce').dt.strftime('%Y-%m-%d')    
    else:
        raise ValueError(f"Parquet 文件中缺少 'date' 列。现有列: {list(df.columns)}")

    # Ensure 'date' column is not an index and is a regular column
    if 'date' in df.index.names:
        df = df.reset_index()

    return convert_bool_fields(df)

def convert_bool_fields(df):
    """Convert boolean fields in DataFrame"""
    bool_cols = {
        'POWERUP': bool,
        'BOTTOMBUY': bool,
        'BOTTOMUPBUY': bool,
        'POTENTIALBUY': bool,
        'STARK': bool,
        'ENTRYBUY': bool,
        'STRONGBUY': bool,
        'TRENDBUY': bool,
        'TRENDSELL': bool,
        'POWERDOWNSELL': bool,
        'CLEANSELL': bool,
        'STAGESELL': bool,
        'DRAGONBUY': bool
    }

    for col, dtype in bool_cols.items():
        if col in df.columns:
            if pd.api.types.is_numeric_dtype(df[col]):
                df[col] = df[col].astype(bool)
            elif pd.api.types.is_object_dtype(df[col]):
                df[col] = df[col].str.lower().map({'true': True, 'false': False})
            elif pd.api.types.is_bool_dtype(df[col]):
                df[col] = df[col].astype(bool)

    return df